import torch
import torch.nn as nn


class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads  # 每个头处理的维度？

        assert (
                self.head_dim * heads == embed_size
        ), "Embedding size needs to be divisible by heads"

        # 定义权重矩阵
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 分割embedding到不同的heads
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)

        # Attention机制的计算
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        #  这里是scaled的self-attention
        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)
        # q 代表 query sequence length（查询序列长度）
        # l: key/value sequence length
        # d: head dimension（每个头的维度大小）
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(
            N, query_len, self.heads * self.head_dim
        )

        # 输出层
        out = self.fc_out(out)
        return out


# 示例使用
embed_size = 18  # 向量深度，越深能够存储的信息越多
heads = 3
attention = SelfAttention(embed_size, heads)
query = values = keys = torch.rand(1, 10, embed_size)  # 示例数据  10个单词，词向量长度为5
mask = None
output = attention(values, keys, query, mask)
print(output.shape)
